/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 * Description: hccl cpu侧对外API接口实现文件
 * Author: gaozhenyeu
 * Create: 2024-04-02
 */

#include "hccl_impl_cpu.h"
#include "hccl_comm_base.h"

namespace hccl {
extern RankTable_t g_hcclDefaultRankTable;

HcclCommBase::HcclCommBase(u64 inCCLbufferSize, u64 outCCLbufferSize, std::string identifier)
    : barrierSendBuf(nullptr), barrierRecvBuf(nullptr), impl_(nullptr),
      inCCLbufferSize_(inCCLbufferSize), outCCLbufferSize_(outCCLbufferSize),
      deviceType_(DevType::DEV_TYPE_COUNT), isFirstBarrier_(true), identifier_(identifier), isHeterogComm_(false),
      lookupRequestPool_(RING_MEMORY_CAPACITY), isSpecialType_(false)
{
    indirectInCCLbuffer_ = DeviceMem();
    indirectOutCCLbuffer_ = DeviceMem();
    barrierInMemory_ = DeviceMem();
    barrierOutMemory_ = DeviceMem();
}

HcclCommBase::~HcclCommBase()
{
    RealeaseBarrierMemory();
    impl_ = nullptr;
}

void HcclCommBase::RealeaseBarrierMemory()
{
    barrierInMemory_.free();
    barrierOutMemory_.free();
}

void HcclCommBase::UpdateIsHaveCpuRank(const RankTable_t &rankTable)
{
    for (u32 i = 0; i < rankTable.rankList.size(); i++) {
        // 同一server的标识IP 是一样的，所以可以以此推算出平均dev个数
        if (rankTable.rankList[i].deviceInfo.devicePhyId == HOST_DEVICE_ID) {
            isHaveCpuRank_ = true;
        }
    }
}

void HcclCommBase::UpdateIsHaveCpuRank(const std::vector<RankInfo> &rankList)
{
    for (u32 i = 0; i < rankList.size(); i++) {
        // 同一server的标识IP 是一样的，所以可以以此推算出平均dev个数
        if (rankList[i].devicePhyId == HOST_DEVICE_ID) {
            isHaveCpuRank_ = true;
        }
    }
}

HcclResult HcclCommBase::InitImpl(DevType deviceType)
{
    return HCCL_SUCCESS;
}

HcclResult HcclCommBase::init(HcclCommParams &params, const RankTable_t &rankTable)
{
    UpdateIsHaveCpuRank(rankTable);
    isHeterogComm_ = params.isHeterogComm;

    HCCL_INFO("hcclComm init workmode [%d]", params.commWorkMode);
    if (params.commWorkMode == WorkMode::HCCL_MODE_AI_CPU) {
        isSpecialType_ = true;
    }

    CHK_RET(InitImpl(params.deviceType));

    /* 强行将最后一个字符置0, 确保其可以做字符串操作 */
    params.id.internal[HCCL_ROOT_INFO_BYTES - 1] = '\0';

    /* 入参判断 */
    if (params.rank >= params.totalRanks) {
        HCCL_ERROR("[HcclComm][Init]errNo[0x%016llx] rank[%u] out of range[0, %u]", HCCL_ERROR_CODE(HCCL_E_PARA),
            params.rank, params.totalRanks - 1);
        return HCCL_E_PARA;
    }
    params.identifier = identifier_;
    CHK_RET(impl_->AtomicInitSet());                  /* 初始化竞争, 只允许被初始化一次 */
    HcclResult ret = impl_->Init(params, rankTable);  /* 初始化实例, 失败则重新开放初始化竞争 */
    if (ret != HCCL_SUCCESS) {
        HCCL_ERROR("[HcclComm][Init]errNo[0x%016llx] hccl initialize failed", HCCL_ERROR_CODE(ret));
        impl_->AtomicInitClear();
        return ret;
    }

    CHK_RET(impl_->InitCCLbuffer(inCCLbufferSize_, outCCLbufferSize_));

    HCCL_USER_CRITICAL_LOG("hcclCommInitInfo:commId[%s], rank[%u], totalRanks[%u], serverId[%s], deviceType[%d]," \
        "logicDevId[%d], identifier[%s]", params.id.internal, params.rank, params.totalRanks, params.serverId.c_str(),
        params.deviceType, params.logicDevId, params.identifier.c_str());
    return HCCL_SUCCESS;
}

HcclResult HcclCommBase::init(HcclCommParams &params, const std::vector<RankInfo> &rankList,
    WorldGroupInfo &groupCommonData)
{
    UpdateIsHaveCpuRank(rankList);
    /* 强行将最后一个字符置0, 确保其可以做字符串操作 */
    params.id.internal[HCCL_ROOT_INFO_BYTES - 1] = '\0';

    HCCL_USER_CRITICAL_LOG("rootInfo[%s], rank[%u], totalRanks[%u], chip[%d], logicDevId[%d]", params.id.internal,
        params.rank, params.totalRanks, params.deviceType, params.logicDevId);

    HCCL_INFO("rootInfo init group workmode[%d]", params.commWorkMode);
    if (params.commWorkMode == WorkMode::HCCL_MODE_AI_CPU) {
        isSpecialType_ = true;
    }

    CHK_RET(InitImpl(params.deviceType));

    /* 入参判断 */
    if (params.rank >= params.totalRanks) {
        HCCL_ERROR("[HcclComm][Init]errNo[0x%016llx] rank[%u] out of range[0, %u]", HCCL_ERROR_CODE(HCCL_E_PARA),
            params.rank, params.totalRanks - 1);
        return HCCL_E_PARA;
    }

    params.identifier = identifier_;
    CHK_RET(impl_->CheckDeviceType(params.deviceType));                /* 芯片类型检查 */
    CHK_RET(impl_->AtomicInitSet());                                 /* 初始化竞争, 只允许被初始化一次 */
    HcclResult ret = impl_->Init(params, rankList, groupCommonData); /* 初始化实例, 失败则重新开放初始化竞争 */
    if (ret != HCCL_SUCCESS) {
        impl_->AtomicInitClear();
        HCCL_ERROR("[HcclComm][Init]errNo[0x%016llx] hccl initialize failed", HCCL_ERROR_CODE(ret));
        return ret;
    }

    return HCCL_SUCCESS;
}

HcclResult HcclCommBase::CreateGroup(const std::string &group, const u32 &groupRank, const u32 &userRank,
    const std::vector<u32> &groupRanks, std::shared_ptr<HcclCommBase> &groupComm)
{
    return HCCL_SUCCESS;
}

HcclResult HcclCommBase::CreateGroup(const std::string &group, const u32 &groupRank, const u32 &userRank,
    const std::vector<u32> &groupRanks, std::shared_ptr<hcclComm> &groupComm)
{
    return HCCL_SUCCESS;
}

HcclResult HcclCommBase::DestroyGroup(const std::string &group) const
{
    /* 增加输出日志关键字 */
    HCCL_DEBUG("start destroy group[%s]", group.c_str());
    return HCCL_SUCCESS;
}

HcclResult HcclCommBase::PsAssociateWorkers(int32_t tag, uint32_t workerRanks[], uint64_t workerNum)
{
    HCCL_INFO("[HcclCommBase]PsAssociateWorkers started.");
    CHK_SMART_PTR_NULL(impl_);
    return impl_->CreateTagAndBuildLinks(tag, workerRanks, workerNum);
}

HcclResult HcclCommBase::DestoryResource(int32_t tag)
{
    HCCL_INFO("[HcclCommBase]DestoryResource started.");
    CHK_SMART_PTR_NULL(impl_);
    return impl_->DestroyCDomainResource(tag);
}

// 目前支持按tag对资源释放、解绑定
HcclResult HcclCommBase::ClearOpResource(const std::string &tag)
{
    CHK_RET(impl_->ClearOpResource(tag));

    return HCCL_SUCCESS;
}

HcclResult HcclCommBase::Break()
{
    impl_->Break();
    return HCCL_SUCCESS;
}

HcclResult HcclCommBase::GetUniqueId(HcclRootInfo *uniqueId)
{
    CHK_PTR_NULL(uniqueId);

    std::string uniqueIdGot = HcclImplBase::GetUniqueId();
    s32 ret = snprintf_s(uniqueId->internal, HCCL_ROOT_INFO_BYTES, HCCL_ROOT_INFO_BYTES - 1,
                         "%s%s", "hccl-", uniqueIdGot.c_str());
    CHK_PRT_RET((ret == -1), HCCL_ERROR("[Get][UniqueId]errNo[0x%016llx] get unique id failed,uniqueId[%p]",
        HCCL_ERROR_CODE(ret), uniqueId), HCCL_E_MEMORY);

    return HCCL_SUCCESS;
}

HcclResult HcclCommBase::CreateCommCCLbuffer() const
{
    CHK_RET(impl_->CreateCommCCLbuffer());

    return HCCL_SUCCESS;
}

std::string HcclCommBase::GetIdentifier()
{
    return identifier_;
}

u32 HcclCommBase::GetRankTableCrc()
{
    return impl_->GetRankTableCrc();
}

HcclResult HcclCommBase::GetUserRank(u32 &userRank)
{
    userRank = impl_->GetUserRank();
    return HCCL_SUCCESS;
}

HcclResult HcclCommBase::GetRankSize(u32 &rankSize)
{
    rankSize = impl_->GetRankSize();

    return HCCL_SUCCESS;
}

HcclResult HcclCommBase::GetWorkspaceSubStreamNum(u64 &streamNum, u64 dataSize, HcclCMDType optype) const
{
    return impl_->GetWorkspaceSubStreamNum(streamNum, dataSize, optype);
}

HcclResult HcclCommBase::GetWorkspaceMemSize(const std::string &opType, u64 count, HcclDataType dataType,
    u32 &rankSize, u64 &size)
{
    return impl_->GetWorkspaceMemSize(opType, count, dataType, rankSize, size, deviceType_);
}

// 设定 workspace 资源
HcclResult HcclCommBase::SetWorkspaceResource(const std::string &tag, void *memPtr, u64 maxSize,
    std::vector<rtStream_t> &stream)
{
    return impl_->SetWorkspaceResource(tag, memPtr, maxSize, stream);
}


HcclResult HcclCommBase::CreateRemoteOpBasedResources(u64 memSize, const std::string &tag)
{
    return impl_->CreateRemoteOpBasedResources(memSize, tag);
}

HcclResult HcclCommBase::CreateOrUpdateRemoteOpBasedResources(u64 memSize, const std::string &tag)
{
    return impl_->CreateOrUpdateRemoteOpBasedResources(memSize, tag);
}

HcclResult HcclCommBase::DestroyRemoteOpBasedMem(const std::string &tag)
{
    return impl_->DestroyRemoteOpBasedMem(tag);
}

HcclResult HcclCommBase::GetHccsLinkNum(u32 &numHccsLink)
{
    return impl_->GetHccsLinkNum(numHccsLink);
}

HcclResult HcclCommBase::GetDevType(DevType &devType)
{
    devType = deviceType_;
    return HCCL_SUCCESS;
}

HcclResult HcclCommBase::IsStandardCard(bool &isStandardCard)
{
    isStandardCard = impl_->IsStandardCard();
    return HCCL_SUCCESS;
}

HcclResult HcclCommBase::CsInit(const RoleTableInfo &roleTableInfo)
{
    CHK_RET(impl_->CsInit(roleTableInfo));
    CHK_RET(lookupRequestPool_.Init());

    return HCCL_SUCCESS;
}

HcclResult HcclCommBase::GetPsMap(std::map<u32, u32> &nodeMap)
{
    CHK_RET(impl_->GetPsMap(nodeMap));
    return HCCL_SUCCESS;
}

HcclOperator<EmbeddingServiceParam> *HcclCommBase::EsOpNew()
{
    return lookupRequestPool_.Alloc();
}

void HcclCommBase::EsOpDelete(HcclOperator<EmbeddingServiceParam> *ptr)
{
    ptr->Clear();
    lookupRequestPool_.Free(ptr);
}

HcclResult HcclCommBase::ServiceCancel(ServiceHandle handle)
{
    auto opPtr = static_cast<HcclOperator<EmbeddingServiceParam> *>(handle);
    opPtr->param_.opType = EmbeddingServiceParam::OperatorType::SERVICE_CANCEL;
    CHK_PRT(impl_->ServiceCancel(opPtr));
    return HCCL_SUCCESS;
}

HcclResult HcclCommBase::SetDevicePid(s32 devicePid)
{
    CHK_PRT(impl_->SetDevicePid(devicePid));
    return HCCL_SUCCESS;
}

HcclResult HcclCommBase::GetTagWorkerList(u32 tag, std::vector<u32> &workerList)
{
    CHK_PRT(impl_->GetTagWorkerList(tag, workerList));
    return HCCL_SUCCESS;
}

HcclResult HcclCommBase::GetWorkerList(std::vector<u32> &workerList)
{
    CHK_PRT(impl_->GetWorkerList(workerList));
    return HCCL_SUCCESS;
}

HcclResult HcclCommBase::GetTopoDesc(HcclTopoDescs *topoDescs, uint32_t topoSize)
{
    HCCL_INFO("HCCL_KEY_INFO: GetTopoDesc topoDescs[%p] topoSize[%u]", topoDescs, topoSize);

    CHK_PRT(impl_->GetTopoDesc(topoDescs, topoSize));

    return HCCL_SUCCESS;
}
}  // namespace hccl
